package com.thinkaurelius.titan.diskstorage.lucene;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;
import com.spatial4j.core.context.SpatialContext;
import com.spatial4j.core.shape.Shape;
import com.thinkaurelius.titan.core.attribute.Geoshape;
import com.thinkaurelius.titan.util.system.IOUtils;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queries.BooleanFilter;
import org.apache.lucene.search.*;
import org.apache.lucene.spatial.SpatialStrategy;
import org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy;
import org.apache.lucene.spatial.prefix.tree.GeohashPrefixTree;
import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
import org.apache.lucene.spatial.query.SpatialArgs;
import org.apache.lucene.spatial.query.SpatialOperation;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;
import org.junit.Before;
import org.junit.Test;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @author Matthias Broecheler (me@matthiasb.com)
 */

public abstract class LuceneExample {

    public static final File path = new File("/tmp/lucene");
    private static final String STR_SUFFIX = "_str";
    private static final String TXT_SUFFIX = "_txt";

    private static final int MAX_RESULT = 10000;

    private Map<String,SpatialStrategy> spatial=new HashMap<String,SpatialStrategy>();
    private SpatialContext ctx = SpatialContext.GEO;

    @Before
    public void setup() {
        if (path.exists()) IOUtils.deleteDirectory(path,false);
        if (!path.exists() && path.isDirectory()) path.mkdirs();
    }

    private SpatialStrategy getSpatialStrategy(String key) {
        SpatialStrategy strategy = spatial.get(key);
        if (strategy==null) {
            final int maxLevels = 11;
            SpatialPrefixTree grid = new GeohashPrefixTree(ctx, maxLevels);
            strategy = new RecursivePrefixTreeStrategy(grid, key);
            spatial.put(key,strategy);
        }
        return strategy;
    }

    @Test
    public void example1() throws Exception {
        Directory dir = FSDirectory.open(path);
        Analyzer analyzer = new StandardAnalyzer();
        IndexWriterConfig iwc = new IndexWriterConfig(Version.LUCENE_4_10_4, analyzer);

        iwc.setOpenMode(IndexWriterConfig.OpenMode.CREATE_OR_APPEND);
        IndexWriter writer = new IndexWriter(dir, iwc);

        indexDocs(writer, "doc1", ImmutableMap.of("name","The laborious work of John Doe as we know it",
                                                  "city","Blumenkamp",
                                                  "location",Geoshape.point(51.687882,6.612053),
                                                  "time",1000342034
                ));

        indexDocs(writer, "doc2", ImmutableMap.of("name","Life as we know it or not",
                "city","Essen",
                "location",Geoshape.point(51.787882,6.712053),
                "time",1000342034-500
        ));

        indexDocs(writer, "doc3", ImmutableMap.of("name","Berlin - poor but sexy and a display of the extraordinary",
                "city","Berlin",
                "location",Geoshape.circle(52.509535, 13.425293, 50),
                "time",1000342034+2000
        ));
        writer.close();

        //Search
        IndexReader reader = DirectoryReader.open(FSDirectory.open(path));
        IndexSearcher searcher = new IndexSearcher(reader);
        analyzer = new StandardAnalyzer();

        //Auesee
        BooleanFilter filter = new BooleanFilter();
        //filter.add(new TermsFilter(new Term("name_txt","know")), BooleanClause.Occur.MUST);

        SpatialArgs args = new SpatialArgs(SpatialOperation.Intersects,Geoshape.circle(51.666167,6.58905,450).convert2Spatial4j());
        //filter.add(getSpatialStrategy("location").makeFilter(args), BooleanClause.Occur.MUST);

        filter.add(NumericRangeFilter.newLongRange("time",(long)1000342034,(long)1000342034,true,true), BooleanClause.Occur.MUST);
//        filter.add(NumericRangeFilter.newLongRange("time",(long)1000342034-100,Long.MAX_VALUE,true,true), BooleanClause.Occur.MUST);
//        filter.add(NumericRangeFilter.newLongRange("time",Long.MIN_VALUE,(long)1000342034+300,true,true), BooleanClause.Occur.MUST);


        filter.add(new PrefixFilter(new Term("city_str","B")), BooleanClause.Occur.MUST);


        TopDocs docs = searcher.search(new MatchAllDocsQuery(), filter, MAX_RESULT);
        if (docs.totalHits>=MAX_RESULT) throw new RuntimeException("Max results exceeded: " + MAX_RESULT);

        Set<String> result = getResults(searcher,docs);
        System.out.println(result);

    }

    private Set<String> getResults(IndexSearcher searcher, TopDocs docs) throws IOException {
        Set<String> found = Sets.newHashSet();
        for (int i = 0; i < docs.totalHits; i++) {
            found.add(searcher.doc(docs.scoreDocs[i].doc).getField("docid").stringValue());
        }
        return found;
    }

    void indexDocs(IndexWriter writer, String docid, Map<String,Object> docMap) throws IOException {
        Document doc = new Document();

        Field docidField = new StringField("docid", docid, Field.Store.YES);
        doc.add(docidField);

        for (Map.Entry<String,Object> kv : docMap.entrySet()) {
            String key = kv.getKey();
            Object value = kv.getValue();

            if (value instanceof Number) {
                Field field = null;
                if (value instanceof Integer || value instanceof Long) {
                    field = new LongField(key, ((Number)value).longValue(), Field.Store.NO);
                } else { //double or float
                    field = new DoubleField(key, ((Number)value).doubleValue(), Field.Store.NO);
                }
                doc.add(field);
            } else if (value instanceof String) {
                String str = (String)value;
                Field field = new TextField(key+ TXT_SUFFIX, str, Field.Store.NO);
                doc.add(field);
                if (str.length()<256)
                field = new StringField(key+STR_SUFFIX, str, Field.Store.NO);
                doc.add(field);
            } else if (value instanceof Geoshape) {
                Shape shape = ((Geoshape)value).convert2Spatial4j();
                for (IndexableField f : getSpatialStrategy(key).createIndexableFields(shape)) {
                    doc.add(f);
                }
            } else throw new IllegalArgumentException("Unsupported type: " + value);
        }

        writer.updateDocument(new Term("docid", docid), doc);

    }


}
